# -*- coding: utf-8 -*-

import logging

from sqlalchemy import select, union, join, asc, Table, MetaData
from sqlalchemy.engine.result import ResultProxy

from py_engine import session, engine
from py_ext.py_declarative import Base, tables

logger = logging.getLogger(__name__)


class User1(Base):
    __table__ = tables['user1']

    pass


class User2(Base):
    __table__ = tables['user2']

    pass


class Role(Base):
    __table__ = tables['role']

    pass


def execute_statement(statement):
    rp: ResultProxy = session.execute(statement)
    logger.info(f'count:{rp.rowcount}')
    batch = 2
    i = 1
    while True:
        l = rp.fetchmany(batch)
        if not l:
            break

        # l: list = rp.fetchall()
        for r in l:
            logger.warning(f'{r.role_name}:{r.user_name}')
            logger.warning(f'{r["role_name"]}:{r["user_name"]}')
            logger.warning(f'{r[0]}:{r[1]}')

        logger.info(f'fetch:{i}')
        i += 1


def py_union1():
    statement1 = select([User1.role_id, User1.user_name1.label('user_name')]).select_from(User1)
    statement2 = select([User2.role_id, User2.user_name2.label('user_name')]).select_from(User2)

    query_list = [statement1, statement2, statement1]
    User = union(*query_list).alias('union_alias')  # noqa

    # statement = select([Role.role_name, User])
    # statement = statement.select_from(join(Role, User, Role.role_id == User.c.role_id))
    # Role.join error
    metadata = MetaData(bind=engine)
    role = Table('role', metadata, autoload=True, autoload_with=engine)
    # statement = statement.select_from(User.join(Role, Role.role_id == User.c.role_id))
    statement = select([role.c.role_name, User.c.user_name]).select_from(
        role.join(User, role.c.role_id == User.c.role_id))
    statement = statement.order_by(asc(Role.role_id))

    execute_statement(statement)


def py_union2():
    statement = """
        select r.role_name, u.user_name from role r
        inner join (
            select role_id, user_name1 user_name  from user1
            union
            select role_id, user_name2 user_name  from user2
        ) u on r.role_id = u.role_id
        order by r.role_id asc
    """
    execute_statement(statement)


def py_union3():
    query1 = session.query(User1.role_id.label('role_id'), User1.user_name1.label('user_name'))
    query2 = session.query(User2.role_id.label('role_id'), User2.user_name2.label('user_name'))
    User = union(query1, query2).alias('union_alias')  # noqa
    l: list = session.query(
        Role.role_name, User
    ).join(
        User, User.c.role_id == Role.role_id
    ).all()
    for r in l:
        logger.warning(f'{r.role_name}:{r.user_name}')
        logger.warning(f'{r[0]}:{r[1]}')
        pass
    pass


if __name__ == '__main__':
    py_union1()
    # py_union2()
    # py_union3()
    pass
